import cv2
import mediapipe as mp

def draw_skeleton_on_video(input_path, output_path):
    mp_pose = mp.solutions.pose
    pose = mp_pose.Pose(
        min_detection_confidence=0.5,
        min_tracking_confidence=0.5,
        model_complexity=2
    )
    
    # 修复1：明确定义所有绘图参数
    connection_style = {
        connection: mp.solutions.drawing_utils.DrawingSpec(
            color=(0, 255, 0),  # BGR颜色格式
            thickness=3,
            circle_radius=None  # 明确禁用关键点绘制
        ) for connection in mp_pose.POSE_CONNECTIONS
    }

    # 修复2：增加视频源验证
    cap = cv2.VideoCapture(input_path)
    if not cap.isOpened():
        print(f"错误：无法打开视频源 {input_path}")
        return

    # 修复3：强制统一颜色通道
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    
    # 修复4：使用更兼容的编码器组合
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, 
                        (frame_width, frame_height),
                        isColor=True)  # 明确指定颜色模式

    try:
        while cap.isOpened():
            success, frame = cap.read()
            if not success:
                break
                
            # 修复5：保持BGR颜色空间一致性
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            results = pose.process(rgb_frame)
            
            if results.pose_landmarks:
                mp.solutions.drawing_utils.draw_landmarks(
                    frame,  # 在原始BGR帧上绘制
                    results.pose_landmarks,
                    mp_pose.POSE_CONNECTIONS,
                    connection_drawing_spec=connection_style
                )
            
            out.write(frame)
    finally:
        # 修复6：确保资源释放
        cap.release()
        out.release()

if __name__ == "__main__":
    # 示例调用（保持英文路径）
    draw_skeleton_on_video("input.mp4", "output.mp4")